"""
dataset_config.py
-----------------
定义与数据集加载相关的配置，包括：
- 数据集类型（内置 / 本地）
- 数据路径、批大小、加载线程数
- 数据拆分比例（train / val / test）
- 是否打乱数据、是否缓存、是否标准化等选项

这些参数主要用于构建 DataLoader。
"""

from dataclasses import dataclass
from typing import Optional, Tuple

from .augment_config import AugmentConfig

# ["MNIST","CIFAR10","Flowers102"]

@dataclass
class DatasetConfig:
    # === 基础信息 ===
    dataset_name: str = "CIFAR10"         # 数据集名称（如 CIFAR10 / IMDB / custom）
    dataset_type: str = "builtin"         # 数据集类型："builtin" 或 "local"
    data_modal: str = "image"             # 数据类型："image" / "text" / "multi-modal"
    data_dir: str = "data"              # 数据路径

    # === 本地数据集信息 ===
    local_dataset_path: Optional[str] = None  # 自定义数据集 Python 文件路径，如 ./datasets/my_dataset.py
    local_dataset_class: Optional[str] = None # 自定义数据集类名，如 MyCustomDataset
    local_train_dir: Optional[str] = None
    local_val_dir: Optional[str] = None
    local_test_dir: Optional[str] = None

    # === DataLoader 参数 ===
    batch_size: int = 64
    num_workers: int = 4
    shuffle: bool = True
    pin_memory: bool = True
    drop_last: bool = False

    # === 数据拆分 ===
    split_ratio: Tuple[float, float, float] = (0.8, 0.1, 0.1)
    """
    数据集拆分比例 (train, val, test)
    当 dataset_type="builtin" 且数据集中无独立验证集时使用。
    三个比例之和应为 1.0
    """

    # === 数据预处理 ===
    normalize: bool = True
    resize: Optional[Tuple[int, int]] = (224, 224)
    cache_data: bool = False

    # === 其他参数 ===
    seed: int = 42

    # 增强配置
    augment: Optional[AugmentConfig] = None

    def get_augment_config(self):
        """
        根据 data_modal 获取对应的增强配置。
        返回 None 时表示没有配置。
        """
        if self.augment is None:
            return None

        if self.data_modal == "image":
            return self.augment.image
        elif self.data_modal == "text":
            return self.augment.text
        else:
            # 对于 multi-modal 或未知类型，可以返回 None 或自定义逻辑
            return None

    def summary(self):
        print("📁 数据集配置")
        print(f"  名称: {self.dataset_name}")
        print(f"  类型: {self.dataset_type}")
        print(f"  数据模态: {self.data_modal}")
        print(f"  数据路径: {self.data_dir}")
        if self.dataset_type == "local":
            print(f"  自定义路径: {self.local_dataset_path}")
            print(f"  类名: {self.local_dataset_class}")
            print(f"  本地数据: train={self.local_train_dir}, val={self.local_val_dir}, test={self.local_test_dir}")
        print(f"  批大小: {self.batch_size}, 工作线程: {self.num_workers}")
        print(f"  拆分比例: {self.split_ratio}")
        print(f"  归一化: {self.normalize}, 调整尺寸: {self.resize}")
        print(f"  缓存数据: {self.cache_data}, 随机种子: {self.seed}")


# -----------------------------
# 测试代码
# -----------------------------
if __name__ == "__main__":
    print("=== 测试 DatasetConfig ===")

    # 图片数据集
    config_image = DatasetConfig(data_modal="image")
    config_image.summary()

    # 文本数据集
    config_text = DatasetConfig(name="IMDB", data_modal="text", batch_size=16)
    config_text.summary()

    # 多模态数据集
    config_multi = DatasetConfig(name="MultiModalDataset", data_modal="multi-modal", batch_size=8)
    config_multi.summary()